#include "stdafx.h"
#include "dasm.h"

namespace ak
{
	CodeParser32::CodeParser32() : _options(0), _imageBase(0), _imageSize(0)
	{
	}

	HRESULT CodeParser32::Init( IPeReader * ope )
	{
		Close();
		if( _ope->IsWow64() ) return E_NOTIMPL;
		_ope = ope;
		return S_OK;
	}


	HRESULT CodeParser32::Close()
	{
		_ope.dispose();
		_options = 0;
		_imageBase = 0;
		_imageSize = 0;
		return S_OK;
	}

	HRESULT CodeParser32::FetchBIL( UINT32 rva, SIZE_T options, TCallback * callback )
	{
		if( !_ope ) return E_UNEXPECTED;
		if( !callback ) return E_INVALIDARG;
		_options = options;
		return Fetch( rva, callback );
	}

	HRESULT CodeParser32::Fetch( UINT32 uRVA, TCallback* callback )
	{
		IPeReader::TImgOptHdrBnd * opt = (IPeReader::TImgOptHdrBnd *)_ope->OptionalHeader();
		if( !opt ) return E_FAIL;
		_imageBase = opt->X32.ImageBase;
		_imageSize = opt->X32.SizeOfImage;
		_stack.push( uRVA );
		while( _stack.size() && SUCCEEDED(DisStackTop(callback) ) );
		return S_OK;
	}

	HRESULT CodeParser32::DisStackTop(TCallback* callback)
	{
		UINT32 uRVA = _stack.top(); _stack.pop();

		RFAILED( callback->OnEvent( DTE_BRANCH ) );

		BOOL bOnce = (_options & DTO_RVA_ONCE);

		for( ;; )
		{
			if( bOnce && (_traced.find(uRVA) != _traced.end()) )
				break;

			_traced.insert(uRVA);

			UINT8 InsBuffer[20] = {};
			DWORD cbRead = _ope->ReadDatRva2( uRVA, InsBuffer, sizeof(InsBuffer) );
			if( !cbRead ) break;

			bptr_t dp( InsBuffer, InsBuffer + cbRead );

			hde32s hid = {};
			size_t il = hde32_disasm( dp, &hid );
			if( !il ) break ;

			if( hid.flags & (F32_ERROR_LENGTH) )
				break ;

			UINT32 uOpcode = 0;
			RFAILED( InsertInstruction32( uRVA, hid, uOpcode, callback ) );

			uRVA += il;

			if( _options & (DTO_JMP|DTO_CALL|DTO_JCC_T) )
			{
				// process branch
				if( !ProcessBranch(uRVA, hid, uOpcode ) )
					break;
			}
		}
		return S_OK;
	}

	HRESULT CodeParser32::InsertInstruction32( UINT32 uRVA, hde32s & hid, UINT32& uOpcode, TCallback* lpCallback )
	{
		UINT16 op16 = hid.opcode;
		if( hid.opcode2 )
		{
			op16 = (op16 << 8) + hid.opcode2;
		}
		else if( (hid.flags & (F32_GROUP|F32_MODRM)) == (F32_GROUP|F32_MODRM) )
		{
			op16 = (op16 << 8) + hid.modrm;
		}

		if( _options & DTO_EX_INS )
		{
			X64INSTEX did;
			did.RVA = uRVA;
			did.InsLen = hid.len;
			did.ModRMRLen = hid.len2;
			did.OpCode = op16;
			did.mix.i32 = hid;
			did.Flag = 0;
			uOpcode = op16;
			return lpCallback->OnInstruction( &did );
		}
		else
		{
			X64INST did;
			did.RVA = uRVA;
			did.InsLen = hid.len;
			did.ModRMRLen = hid.len2;
			did.OpCode = op16;
			did.Flag = 0;
			uOpcode = op16;
			return lpCallback->OnInstruction( &did );
		}
	}

	VOID CodeParser32::TraceRVA( UINT32 uRVA )
	{
		_stack.push( uRVA );
	}

	BOOL CodeParser32::ProcessBranch( UINT32& uRVA, hde32s& hid, UINT32 uOpcode )
	{
		if( _options & DTO_ZERO_BREAK )
		{
			if( 0 == uOpcode && 0 == hid.modrm && 0 == hid.modrm_mod && 0 == hid.modrm_reg && 0 == hid.modrm_rm)
				return FALSE;
		}

		if( _options & DTO_CALL )
		{
			if( uOpcode == 0xE8 )
			{
				INT32 off = hid.imm.imm32;
				if( _options & DTO_BRACH_FIRST )
				{
					TraceRVA( uRVA );
					uRVA += off;
					if( uRVA >= _imageSize ) return FALSE;
				}
				else
				{
					TraceRVA( uRVA + off );
				}
				return TRUE;
			}

			if( uOpcode == 0xC2 || uOpcode == 0xC3 )
				return FALSE;
		}

		if( _options & DTO_JMP )
		{
			if( uOpcode == 0xEB )
			{
				INT8 off = hid.imm.imm8;
				uRVA += off;
				if( uRVA >= _imageSize ) return FALSE;
				return TRUE;
			}
			else if( uOpcode == 0xE9 )
			{
				INT32 off = hid.imm.imm32;
				uRVA += off;
				if( uRVA >= _imageSize ) return FALSE;
				return TRUE;
			}
			else if( uOpcode == 0xFF25 )
			{
				UINT32 uReadRVA = hid.disp.disp32;
				if( uReadRVA < _imageBase ) return FALSE;
				uReadRVA -= _imageBase;
				if( uReadRVA >= _imageSize ) return FALSE;
				// FIX ME 
				return FALSE;
				//if( SUCCEEDED(cp32->GetSymbolByIAT(uReadRVA,NULL,FALSE) ) )
				//	return FALSE;
				le::U32 uVA;
				if( !_ope->ReadDatRva( uReadRVA, &uVA, sizeof(uVA) ) )
					return FALSE;
				if( uVA < _imageBase ) return FALSE;
				uReadRVA = uVA;
				uReadRVA -= _imageBase;
				if( uReadRVA >= _imageSize ) return FALSE;
				uRVA = uReadRVA;
				return TRUE;
			}
			SIZE_T tmp = uOpcode & 0xFFF0;
			if( tmp == 0xFF20 || tmp == 0xFF60 || tmp == 0xFFA0 || tmp == 0xFFE0 )
			{
				return FALSE;
			}
		}

		if( _options & DTO_JCC_T )
		{
			if( (uOpcode & 0xFFF0) == 0x70 )
			{
				INT8 off = hid.imm.imm8;
				if( _options & DTO_BRACH_FIRST )
				{
					TraceRVA( uRVA );
					uRVA += off;
					return TRUE;
				}
				TraceRVA( uRVA + off );
				return TRUE;
			}
			if( (uOpcode & 0xFFF0) == 0x0F80 )
			{
				INT32 off = hid.imm.imm32;
				if( _options & DTO_BRACH_FIRST )
				{
					TraceRVA( uRVA );
					uRVA += off;
					return TRUE;
				}
				TraceRVA( uRVA + off );
				return TRUE;
			}
		}
		return TRUE;
	}
};
